{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b63df770-ff68-4212-8c0e-2be0380f52d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/miniconda3/envs/py310/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a22cf280-b9ab-4bd0-81c6-5a350ffcf3ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = '/model/ModelScope/Qwen/Qwen3-0.6B'\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id, device_map=\"auto\", torch_dtype=\"auto\", trust_remote_code=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7f509dc3-e10f-4830-928d-d4bd956437a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Qwen3ForCausalLM(\n",
       "  (model): Qwen3Model(\n",
       "    (embed_tokens): Embedding(151936, 1024)\n",
       "    (layers): ModuleList(\n",
       "      (0-27): 28 x Qwen3DecoderLayer(\n",
       "        (self_attn): Qwen3Attention(\n",
       "          (q_proj): Linear(in_features=1024, out_features=2048, bias=False)\n",
       "          (k_proj): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "          (v_proj): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "          (o_proj): Linear(in_features=2048, out_features=1024, bias=False)\n",
       "          (q_norm): Qwen3RMSNorm((128,), eps=1e-06)\n",
       "          (k_norm): Qwen3RMSNorm((128,), eps=1e-06)\n",
       "        )\n",
       "        (mlp): Qwen3MLP(\n",
       "          (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)\n",
       "          (up_proj): Linear(in_features=1024, out_features=3072, bias=False)\n",
       "          (down_proj): Linear(in_features=3072, out_features=1024, bias=False)\n",
       "          (act_fn): SiLU()\n",
       "        )\n",
       "        (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)\n",
       "        (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)\n",
       "      )\n",
       "    )\n",
       "    (norm): Qwen3RMSNorm((1024,), eps=1e-06)\n",
       "    (rotary_emb): Qwen3RotaryEmbedding()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=1024, out_features=151936, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "47df61fd-5b7c-4ee0-87d6-afd9374a4054",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(151936, 1024)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_origin = model.model.embed_tokens.weight.cpu().detach().float().numpy()\n",
    "embedding_origin.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7d036a0d-9dc1-4738-815d-3a92b98261c4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2048, 1024)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "weights_origin = model.model.layers[0].self_attn.q_proj.weight.cpu().detach().float().numpy()\n",
    "rows_origin, cols_origin = weights_origin.shape\n",
    "rows_origin, cols_origin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3eff032a-314a-420b-90b4-0814c197e4cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# strip = 8\n",
    "# weights = weights_origin[::strip, ::strip]\n",
    "\n",
    "sub_mat_size = 8\n",
    "\n",
    "weights = weights_origin.reshape(rows_origin//sub_mat_size, sub_mat_size, cols_origin//sub_mat_size, sub_mat_size).swapaxes(1, 2).reshape(-1, sub_mat_size, sub_mat_size).max(axis=(1, 2)).reshape(rows_origin//sub_mat_size,cols_origin//sub_mat_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a16ba0bf-6c04-427b-a130-4f6ec616e13c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.64453125"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights.max().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "cbd600c7-eefa-4511-aae3-e6888c535b60",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[  0   1   2 ... 253 254 255]\n",
      " [  0   1   2 ... 253 254 255]\n",
      " [  0   1   2 ... 253 254 255]\n",
      " ...\n",
      " [  0   1   2 ... 253 254 255]\n",
      " [  0   1   2 ... 253 254 255]\n",
      " [  0   1   2 ... 253 254 255]] [[  0   0   0 ...   0   0   0]\n",
      " [  1   1   1 ...   1   1   1]\n",
      " [  2   2   2 ...   2   2   2]\n",
      " ...\n",
      " [509 509 509 ... 509 509 509]\n",
      " [510 510 510 ... 510 510 510]\n",
      " [511 511 511 ... 511 511 511]]\n",
      "[[0. 0. 0. ... 0. 0. 0.]\n",
      " [0. 0. 0. ... 0. 0. 0.]\n",
      " [0. 0. 0. ... 0. 0. 0.]\n",
      " ...\n",
      " [0. 0. 0. ... 0. 0. 0.]\n",
      " [0. 0. 0. ... 0. 0. 0.]\n",
      " [0. 0. 0. ... 0. 0. 0.]]\n",
      "执行时间：186.351504 秒\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib\n",
    "import random\n",
    "\n",
    "matplotlib.use('Agg')\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "plt.ioff()\n",
    "\n",
    "plt.rcParams['path.simplify'] = True\n",
    "plt.rcParams['path.simplify_threshold'] = 0.1\n",
    "\n",
    "# 创建行列索引网格\n",
    "rows, cols = weights.shape\n",
    "\n",
    "x = np.arange(0, cols)  # 列索引作为x轴\n",
    "y = np.arange(0, rows)  # 行索引作为y轴\n",
    "\n",
    "X, Y = np.meshgrid(x, y)  # 生成坐标网格\n",
    "\n",
    "print(X, Y)\n",
    "\n",
    "# 创建三维图形\n",
    "fig = plt.figure(figsize=(10, 8))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "\n",
    "bottom = np.zeros_like(weights)\n",
    "print(bottom)\n",
    "\n",
    "# colors = [\"blue\", \"cornflowerblue\", \"mediumturquoise\", \"goldenrod\"] * int(rows * cols / 4)\n",
    "\n",
    "colors_element=[\"red\", \"yellow\", \"blue\", \"green\", \"orange\"]\n",
    "colors=[]\n",
    "temp = random.choices(colors_element, k=cols)\n",
    "for i in range(0, rows):\n",
    "    colors.extend(temp)\n",
    "\n",
    "colors_np =np.array(colors)\n",
    "# print(colors_np)\n",
    "\n",
    "ax.bar3d(\n",
    "    X.ravel(), Y.ravel(), bottom.ravel(), 0.01, 0.01, np.abs(weights).ravel(),\n",
    "    color=colors_np,\n",
    "    shade=True,\n",
    ")\n",
    "\n",
    "# 添加标签和标题\n",
    "ax.set_xlabel('in feature', labelpad=15)\n",
    "ax.set_ylabel('out feature', labelpad=15)\n",
    "ax.set_zlabel('Value', labelpad=10)\n",
    "ax.set_title('Linear Layer Tensor', pad=20)\n",
    "ax.set_zlim(0, weights.max().item())\n",
    "\n",
    "start = time.perf_counter()\n",
    "plt.savefig(\"demo_strip8.pdf\", dpi=100)\n",
    "end = time.perf_counter()\n",
    "exe_time = end - start\n",
    "print(f\"执行时间：{exe_time:.6f} 秒\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d31cbf79-ca5c-4d60-a51d-684124f4cf06",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
